# -*- coding: utf-8 -*-
# time: 2025/4/2 16:40
# file: min_model.py
# author: hanson
# 3. 初始化模型
from langchain_chroma import Chroma
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_ollama import OllamaEmbeddings, ChatOllama

# 模型名称
MODEL_NAME = "deepseek-r1:7b"
embeddings = OllamaEmbeddings(model=MODEL_NAME)
llm = ChatOllama(model=MODEL_NAME, temperature=0.3)
VECTOR_DIR = "data/vector_store"
# 4. 创建向量数据库（自动持久化）
# 清理旧向量数据库的数据

"""
vector_db = Chroma.from_documents(
    documents=split_docs,
    embedding=embeddings,
    persist_directory=VECTOR_DIR,
    collection_name="latest_knowledge",
)
"""
print("向量数据库已创建")

# 5. 构建检索链（优化中文模板）

prompt_template = """
您是一个设计用于査询文档来回答问题的代理。您可以使用文档检索工具。
并基于检索内容来回答问题您可能不查询文档就知道答案，但是您仍然应该查询文档来获得答案。
如果您从文档中找不到任何信息用于回答问题，则只需返回“抱歉，这个问题我还不知道。”作为答案。
如果有人提问等关于您的名字的问题，您就回答：“我是超级牛逼哄哄的小天才助手”作为答案。
上下文：{context}
问题：{question}
"""

# 1. 初始化向量数据库
vector_store = Chroma(
    persist_directory=VECTOR_DIR,
    embedding_function=embeddings,
)

# 3. 初始化检索器，并设置检索参数
retriever = vector_store.as_retriever(
    search_type="mmr",
    search_kwargs={
        "k": 2,
        "fetch_k": 20,
        "lambda_mult": 0.5,
        "score_threshold": 0.4,
    },
)

chain = (
    {"context": retriever, "question": RunnablePassthrough()}
    | ChatPromptTemplate.from_template(prompt_template)
    | llm
)

# 6. 执行查询
response = chain.invoke("为什么我的开发者认证状态显示为已取消？")
print(f"\n回答：\n{response.content}")
